from .dataset import FundusDataset
from torch.utils.data import DataLoader
from torchvision import transforms
import lightning as L
class FundusDataModule(L.LightningDataModule):
    def __init__(self, data_dir, batch_size=32, input_size=224, num_workers=4, dataset_name='GlaucomaFundus', split_name='official', use_val=False):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.input_size = input_size
        self.mean = (0.4237, 0.2609, 0.1284)
        self.std = (0.2948, 0.2016, 0.1366)  # from retfound
        self.num_workers = num_workers
        self.dataset_name = dataset_name
        self.split_name = split_name
        self.use_val = use_val  # 更正为一致的命名
        self.val_dataset = None  # 初始化验证集为 None

    def on_train_start(self):
        activate_device = self.device  # 获取当前设备
        self.logger.info(f"Using device {activate_device}")  # 记录到 logger
        print(f"Using device {activate_device}")  # 控制台打印

    def setup(self, stage=None):
        # 定义训练时的变换
        train_transform = transforms.Compose([
            transforms.RandomResizedCrop(self.input_size),
            transforms.RandomHorizontalFlip(),
            transforms.RandomVerticalFlip(),
            transforms.ToTensor(),
            transforms.Normalize(self.mean, self.std),
        ])

        # 定义测试/验证时的变换
        test_transform = transforms.Compose([
            transforms.Resize((self.input_size, self.input_size)),
            transforms.ToTensor(),
            transforms.Normalize(self.mean, self.std),
        ])

        # 创建训练和测试数据集
        self.train_dataset = FundusDataset(root=self.data_dir, split='train', split_name=self.split_name, transform=train_transform, dataset_name=self.dataset_name)
        self.test_dataset = FundusDataset(root=self.data_dir, split='test', split_name=self.split_name, transform=test_transform, dataset_name=self.dataset_name)

        # 如果 use_val 为 True，则创建验证数据集
        if self.use_val:
            self.val_dataset = FundusDataset(root=self.data_dir, split='val', split_name=self.split_name, transform=test_transform, dataset_name=self.dataset_name)

    def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers)

    def test_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers)

    def predict_dataloader(self):
        return DataLoader(self.test_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers)

    def val_dataloader(self):
        if self.use_val:
            return DataLoader(self.val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=self.num_workers)
        else:
            return None  # 不使用验证集时返回 None